
#include"ndtSlam.h"
#include"utils.h"
#include"register.h"

#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/search/impl/search.hpp>
#include <pcl/range_image/range_image.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/common/common.h>
#include <pcl/common/transforms.h>
#include <pcl/registration/icp.h>
#include <pcl/io/pcd_io.h>
#include <pcl/io/ply_io.h>
#include <pcl/filters/filter.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/crop_box.h> 
#include <pcl_conversions/pcl_conversions.h>

#include <tf/transform_datatypes.h>
#include <tf/transform_broadcaster.h>



pcl::PointCloud<PT_TYPE> transformPointCloud(pcl::PointCloud<PT_TYPE>& cloudIn, PointTypePose *transformIn)
{
    pcl::PointCloud<PT_TYPE> cloudOut;
    PT_TYPE* pointFrom;
    int cloudSize = cloudIn.size();
    cloudOut.resize(cloudSize);

    Eigen::Affine3f transCur = pcl::getTransformation(transformIn->x, transformIn->y, transformIn->z, transformIn->roll, transformIn->pitch, transformIn->yaw);
    //#pragma omp parallel for num_threads(numberOfCores)
    for (int i = 0; i < cloudSize; ++i)
    {
        pointFrom = &cloudIn.points[i];
        cloudOut.points[i].x = transCur(0, 0) * pointFrom->x + transCur(0, 1) * pointFrom->y + transCur(0, 2) * pointFrom->z + transCur(0, 3);
        cloudOut.points[i].y = transCur(1, 0) * pointFrom->x + transCur(1, 1) * pointFrom->y + transCur(1, 2) * pointFrom->z + transCur(1, 3);
        cloudOut.points[i].z = transCur(2, 0) * pointFrom->x + transCur(2, 1) * pointFrom->y + transCur(2, 2) * pointFrom->z + transCur(2, 3);
        //cloudOut.points[i].intensity = pointFrom->intensity;
    }

    return cloudOut;
}

Eigen::Affine3f pclPointToAffine3f(PointTypePose thisPoint)
{
    return pcl::getTransformation(thisPoint.x, thisPoint.y, thisPoint.z, thisPoint.roll, thisPoint.pitch, thisPoint.yaw);
}


CNdtSlam::CNdtSlam()
{
    int num_threads = omp_get_max_threads();
    printf("max threads: %d \n", num_threads);
    mOmpNdt.setResolution(1.0);
    mOmpNdt.setNumThreads(num_threads-1);
    mOmpNdt.setNeighborhoodSearchMethod(pclomp::DIRECT7);

    surroundingkeyframeAddingDistThreshold  = 0.1  ;  //# meters, regulate keyframe adding threshold
    surroundingkeyframeAddingAngleThreshold = 0.01 ;  //# radians, regulate keyframe adding threshold
}

CNdtSlam::~CNdtSlam()
{

}

bool CNdtSlam::SaveFrame()
{
    if (mKeyFramesPose6D.points.empty())
        return true;

    Eigen::Affine3f transStart = pclPointToAffine3f(mKeyFramesPose6D.back());
    Eigen::Affine3f transFinal = pclPointToAffine3f(mCurrentPose);
    Eigen::Affine3f transBetween = transStart.inverse() * transFinal;
    float x, y, z, roll, pitch, yaw;
    pcl::getTranslationAndEulerAngles(transBetween, x, y, z, roll, pitch, yaw);

    if (abs(roll)  < surroundingkeyframeAddingAngleThreshold &&
        abs(pitch) < surroundingkeyframeAddingAngleThreshold &&
        abs(yaw)   < surroundingkeyframeAddingAngleThreshold &&
        sqrt(x * x + y * y + z * z) < surroundingkeyframeAddingDistThreshold )
        return false;

    return true;
}


int CNdtSlam::AddImu(ImuData& imu)
{
    mLidarUndistort.AddImu(imu);

    return 0;
}

int CNdtSlam::AddScan(LidarFrame& scan)
{
    pcl::PLYWriter ply_writer;

    //对点云进行畸变校正
    // pcl::PointCloud<PT_TYPE> cloudDistorted;
    // for(int i=0; i<scan.rawCloud.size(); i++)
    // {
    //     PT_TYPE pt;
    //     pt.x = scan.rawCloud[i].x;
    //     pt.y = scan.rawCloud[i].y;
    //     pt.z = scan.rawCloud[i].z;
    //     cloudDistorted.push_back(pt);
    // }
    // ply_writer.write("cloud-distorted.ply", cloudDistorted);
    int res = mLidarUndistort.Process( scan.rawCloud, scan.stamp );
    if(res<0)
        return -1;

    pcl::PointCloud<PT_TYPE> cloud;
    for(int i=0; i<scan.rawCloud.size(); i++)
    {
        PT_TYPE pt;
        pt.x = scan.rawCloud[i].x;
        pt.y = scan.rawCloud[i].y;
        pt.z = scan.rawCloud[i].z;
        cloud.push_back(pt);
    }
    //ply_writer.write("cloud-undistorted.ply", cloud);

    VoxelFilter(cloud, 0.2);

    if(mKeyFrames.size()==0)
    {
        mKeyFrames.push_back(cloud);
        
        PointXYZI position;
        position.x = 0;
        position.y = 0;
        position.z = 0;
        mKeyFramesPose3D.push_back(position);

        PointTypePose pos;
        pos.x = 0;
        pos.y = 0;
        pos.z = 0;
        pos.roll = 0;
        pos.pitch = 0;
        pos.yaw = 0;
        mKeyFramesPose6D.push_back(pos);

        mCurrentPose = pos;
        mCurrentStamp = scan.stamp;

        return 0;
    }

    //如果有IMU数据，就根据IMU来预测相对的姿态变化
    Eigen::Affine3f increRot = mLidarUndistort.PredictIncreRot(mCurrentStamp, scan.stamp);

    //根据相对变化和上一帧的位置来计算新输入帧的位置和姿态
    Eigen::Affine3f currentPose = pcl::getTransformation(
        mCurrentPose.x, mCurrentPose.y, mCurrentPose.z, 
        mCurrentPose.roll, mCurrentPose.pitch, mCurrentPose.yaw);
    Eigen::Affine3f  scanPosePred = currentPose*increRot;
    
    //find the keyframes using radius 
    // std::vector<int>   pointSearchInd;
    // std::vector<float> pointSearchSqDis;
    // pcl::KdTreeFLANN<PointXYZI> kdtree; 
    // kdtree.setInputCloud(mKeyFramesPose3D.makeShared());
    // kdtree.radiusSearch(mKeyFramesPose3D.back(), 5, pointSearchInd, pointSearchSqDis);

    pcl::PointCloud<PT_TYPE> localMap; 
    //printf("key frame size: %d \n", mKeyFrames.size());
    int lindex = mKeyFrames.size()-5;
    if (lindex<0) lindex=0;
    for(int i=lindex; i<mKeyFrames.size(); i++)
    {
        //printf("%d ", i);
        if(i<0) continue;
        //int index = pointSearchInd[i];       
        localMap += transformPointCloud( mKeyFrames[i], &mKeyFramesPose6D.points[i]) ;
    }
    //printf("local map size: %d \n", localMap.size());
    VoxelFilter(localMap, 0.2);

    
    mOmpNdt.setInputTarget(localMap.makeShared());
    mOmpNdt.setInputSource(cloud.makeShared());
    pcl::PointCloud<PT_TYPE>::Ptr aligned(new pcl::PointCloud<PT_TYPE>());
    //auto t1 = ros::WallTime::now();
    mOmpNdt.align(*aligned, scanPosePred.matrix());
    //auto t2 = ros::WallTime::now();
    //std::cout << "single : " << (t2 - t1).toSec() * 1000 << "[msec]" << std::endl;
    float score = mOmpNdt.getFitnessScore();
    Eigen::Matrix4f mt = mOmpNdt.getFinalTransformation();
    
    
    // Eigen::Matrix4d mt;
    // gicpMatch(cloud, localMap, mt); 

    
    // ply_writer.write("ndt-local.ply", localMap);
    // ply_writer.write("ndt-scan.ply", cloud);
    // //ply_writer.write("ndt-align.ply", *aligned);
    // pcl::PointCloud<PT_TYPE> predictCloud;
    // pcl::transformPointCloud(cloud, predictCloud, scanPosePred.matrix());
    // ply_writer.write("ndt-predict.ply", predictCloud);

    //保存当前帧的位置和姿态
    Eigen::Quaternionf mq( mt.block<3,3>(0,0) );
    tf::Quaternion orientation( mq.x(), mq.y(), mq.z(), mq.w() );
    double roll,pitch,yaw;
    tf::Matrix3x3(orientation).getRPY(roll, pitch, yaw);
    PointTypePose scanPose;
    scanPose.x = mt(0,3);
    scanPose.y = mt(1,3);
    scanPose.z = mt(2,3);
    scanPose.roll  = roll;
    scanPose.pitch = pitch;
    scanPose.yaw   = yaw;

    mCurrentPose = scanPose;
    mCurrentStamp = scan.stamp;

    if(SaveFrame())
    {
        printf("xyz: %f %f %f   rpy: %f %f %f \n", scanPose.x,scanPose.y,scanPose.z,
            scanPose.roll, scanPose.pitch, scanPose.yaw);

        PointXYZI position;
        position.x = scanPose.x;
        position.y = scanPose.y;
        position.z = scanPose.z;
        mKeyFramesPose3D.push_back(position);
        mKeyFramesPose6D.push_back(scanPose);
        mKeyFrames.push_back(cloud);
    }

    return 0;
}

int  CNdtSlam::OutputMap(string file)
{
    pcl::PointCloud<PT_TYPE> localMap; 
    for(int i=0; i<mKeyFrames.size(); i++)
    {
        localMap += transformPointCloud( mKeyFrames[i], &mKeyFramesPose6D.points[i]) ;
    }
    VoxelFilter(localMap, 0.2);

    pcl::PLYWriter ply_writer;
    ply_writer.write(file, localMap);

    return 0;
}